#!/usr/bin/env python
'''This runs a sequence of commands on a remote host using SSH. It runs a
simple system checks such as uptime and free to monitor the state of the remote
host.
It works like this:
    Login via SSH (This is the hardest part).
    Run and parse 'uptime'.
    Run 'iostat'.
    Run 'vmstat'.
    Run 'netstat'
    Run 'free'.
    Exit the remote host

1. show all avilable hosts
    xmon -l

2. specify the remote host name
    xmon host1

3. specify the host name pattern, select from many hosts
    xmon host[1-2]

4. input username and password in command line, it will try default config firstlly
    xmon 192.168.0.100#22

5. exec a command sequence after default commands
    xmon host1 -a "ps aux | grep mysql"

6. exec a command sequence instead of default commands, use ',' as delimiter
    xmon host1 -d, -c "ls, free -m"
'''

import re
import os
import sys
import optparse
import traceback
import getpass
import time

from pexpect import *
from parse_ini import *

def search_host(pat, xtools_hosts):
    ''' search host from pat or user\'s selection '''
    def mask_password(password):
        ''' protect the password '''
        length = len(password)
        if length < 8:
            return '*' * length
        else:
            return password[:3] + '*' * (length-6) + password[-3:]

    if pat.find('#') != -1:
        host_pat, port = pat.split('#', 1)
    else:
        host_pat, port = pat, None

    select_hosts = {}
    for name, host in xtools_hosts.items():
        if re.search(host_pat, name):
            select_hosts[name] = host

    nr = len(select_hosts)
    if nr == 0:
        common_config = xtools_hosts.get(COMMON_SECTION_NAME, {})
        username = common_config.get('username')
        password = common_config.get('password')
        if port is None:
            port = common_config.get('port')
        print 'connect to %s@%s#%s using password:%s' % (username, host_pat, port, mask_password(password))
        return {'host':host_pat, 'username':username, 'password':password, 'port':port}    
    elif nr == 1:
        return select_hosts.values()[0]
    else:
        show_available_hosts(select_hosts)
        name = raw_input('which name? ')
        return select_hosts.get(name)
    
    
def try_login(remote_host, try_nr = 1):
    if try_nr > 3:
        print 'You have tried 3 times, now exit'
        return
    else:    
        try_nr += 1
    
    # build ssh command
    ssh_options = '-q'
    port = remote_host.get('port')
    if port:
        ssh_options += ' -p ' + str(port)    
    cmd = 'ssh %s -l %s %s' % (ssh_options, remote_host['username'], remote_host['host'])
    
    # This does not distinguish between a remote server 'password' prompt
    # and a local ssh 'passphrase' prompt (for unlocking a private key).
    handle = spawn(cmd)
    events = ["(?i)are you sure you want to continue connecting", 
              "(?i)(?:password)|(?:passphrase for key)", 
              "(?i)permission denied", 
              "(?i)terminal type", 
              "(?i)connection closed by remote host",
              r'[#>\]$]', 
              EOF,
              TIMEOUT]
    i = handle.expect(events, timeout=2)

    # First phase
    if i==0: 
        # New certificate -- always accept it.
        # This is what you get if SSH does not have the remote host's
        # public key stored in the 'known_hosts' cache.
        handle.sendline("yes")
        i = handle.expect(events, timeout=2)
    if i==1: # password or passphrase
        handle.sendline(remote_host['password'])
        i = handle.expect(events)

    # Second phase
    if i==0:
        # This is weird. This should not happen twice in a row.
        handle.close()
        raise ExceptionPexpect ('Weird error. Got "are you sure" prompt twice.')
    elif i==1: # password prompt again
        # For incorrect passwords, some ssh servers will
        # ask for the password again, others return 'denied' right away.
        # If we get the password prompt again then this means
        # we didn't get the password right the first time.
        print '%s login failed, password refused' % remote_host['username']
        remote_host['username'] = raw_input('username for %s: ' % HOST_PATTERN)
        remote_host['password'] = getpass.getpass('%s\'s password: ' % remote_host['username'])        
        handle.close()
        handle = try_login(remote_host, try_nr)        
        if handle:
            record_password(remote_host, XTOOLS_HOSTS)
    elif i==2: # permission denied -- password was bad.
        handle.close()
        print 'permission denied for %s@%s' % (remote_host['username'], remote_host['password'])
        return
    elif i==3: # terminal type again? WTF?
        handle.close()
        raise ExceptionPexpect ('Weird error. Got "terminal type" prompt twice.')
    elif i==4: # Connection closed by remote host
        handle.close()
        raise ExceptionPexpect ('connection closed')
    elif i==5: # can occur if you have a public key pair set to authenticate. 
        ### TODO: May NOT be OK if expect() got tricked and matched a false prompt.
        pass
    elif i==6:
        handle.close()
        print '%s@%s#%s login failed' % (remote_host['username'], remote_host['host'], remote_host['port'])
        return
    elif i==7: # Timeout
        #This is tricky... I presume that we are at the command-line prompt.
        #It may be that the shell prompt was so weird that we couldn't match
        #it. Or it may be that we couldn't log in for some other reason. I
        #can't be sure, but it's safe to guess that we did login because if
        #I presume wrong and we are not logged in then this should be caught
        #later when I try to set the shell prompt.
        handle.close()
        print 'timeout when connecting ', remote_host['host']
        return
    else: # Unexpected 
        handle.close()
        raise ExceptionPexpect ('unexpected login response')
    
    return handle    

def synch_original_prompt (handle):
    """This attempts to find the prompt. Basically, press enter and record
    the response; press enter again and record the response; if the two
    responses are similar then assume we are at the original prompt. """

    # All of these timing pace values are magic.
    # I came up with these based on what seemed reliable for
    # connecting to a heavily loaded machine I have.
    # If latency is worse than these values then this will fail.

    def levenshtein_distance(a,b):
        """This calculates the Levenshtein distance between a and b.
        """

        n, m = len(a), len(b)
        if n > m:
            a,b = b,a
            n,m = m,n
        current = range(n+1)
        for i in range(1,m+1):
            previous, current = current, [i]+[0]*n
            for j in range(1,n+1):
                add, delete = previous[j]+1, current[j-1]+1
                change = previous[j-1]
                if a[j-1] != b[i-1]:
                    change = change + 1
                current[j] = min(add, delete, change)
        return current[n]

    handle.sendline('\n\n\n')
    time.sleep(0.2)
    out = handle.read_nonblocking(size=8192, timeout=1)
    if out.find('\r\n') != -1:
       sep = '\r\n'
    elif out.find('\n') != -1:
       sep = '\n'
    else:
       sep = '\r'
    lines = [line for line in out.split(sep) if len(line.strip()) > 0]    
    a, b = lines[-2:]
    ld = levenshtein_distance(a,b)
    len_a = len(a)
    if float(ld)/len_a < 0.4:
       return True
    return False



def set_unique_prompt (handle, sh_prompt, set_sh_prompt, set_csh_prompt):
    """This sets the remote prompt to something more unique than # or $.
    This makes it easier for the prompt() method to match the shell prompt
    unambiguously. This method is called automatically by the login()
    method, but you may want to call it manually if you somehow reset the
    shell prompt. For example, if you 'su' to a different user then you
    will need to manually reset the prompt. This sends shell commands to
    the remote host to set the prompt, so this assumes the remote host is
    ready to receive commands.

    Alternatively, you may use your own prompt pattern. Just set the PROMPT
    attribute to a regular expression that matches it. In this case you
    should call login() with auto_prompt_reset=False; then set the PROMPT
    attribute. After that the prompt() method will try to match your prompt
    pattern."""

    handle.sendline ("unset PROMPT_COMMAND")
    handle.sendline (set_sh_prompt) # sh-style
    i = handle.expect ([TIMEOUT, sh_prompt], timeout=1)
    if i == 0: # csh-style
        handle.sendline (set_csh_prompt)
        i = handle.expect ([TIMEOUT, sh_prompt], timeout=1)
        if i == 0:
            return False
    return True


def monitor(handle, shell_prompt):
    ''' send a sequnence of commands and parse their outputs '''

    if OPTIONS.commands:
        commands = OPTIONS.commands.split(OPTIONS.delimiter)
    else:
        commands = ['uptime',
                    'iostat', 
                    'vmstat', 
                    'df -h', 
                    'free -m', 
                    'netstat -lntu', 
                    'ps -ef | grep -c httpd',
                    #'lsof -i4',
                    ]
        if OPTIONS.append:
            commands.extend(OPTIONS.append.split(OPTIONS.delimiter))

    # Now we should be at the command prompt and ready to run some commands.
    for cmd in commands:
        print '/' + '=' * 79
        print '|', cmd 
        print '\\' + '-' * 79
        handle.sendline(cmd)
        handle.expect(shell_prompt)
        print handle.before.split('\n', 1)[1]

    # Now exit the remote host.
    handle.sendline ('exit')
    index = handle.expect([EOF, "(?i)there are stopped jobs"])
    if index==1:
        handle.sendline("exit")
        handle.expect(EOF)


def main():
    remote_host = search_host(HOST_PATTERN, XTOOLS_HOSTS)
    if not remote_host:
        print 'no host was chosen'
        return
    
    handle = try_login(remote_host)
    if not handle:
        print '%-60s - %s' % ('login ' + remote_host['host'], 'FAIL')
        return
    else:
        print '%-60s - %s' % ('login ' + remote_host['host'], 'OK')
    #
    # Set command prompt to something more unique.
    #
    if not synch_original_prompt(handle):
        handle.close()
        raise ExceptionPexpect('could not synchronize with original prompt')
    else:
        print '%-60s - %s' % ('synchronize with original prompt', 'OK')

    shell_prompt = r"\[PEXPECT\][$#] "
    set_sh_prompt = "PS1='[PEXPECT]\$ '" 
    set_csh_prompt = "set prompt='[PEXPECT]\$ '"
    if not set_unique_prompt(handle, shell_prompt, set_sh_prompt, set_csh_prompt):
        handle.close()
        raise ExceptionPexpect('could not set shell prompt\n'+handle.before)
    else:
        print '%-60s - %s' % ('set unique shell prompt', 'OK')

    monitor(handle, shell_prompt)
 
        
if __name__ == '__main__':
    try:
        PARSER = optparse.OptionParser(formatter=optparse.TitledHelpFormatter(), 
                usage=globals()['__doc__'], 
                version='1.0 2009-07-13 colinli',
                conflict_handler="resolve")
        PARSER.add_option ('-v', '--verbose', action='store_true', default=False, help='verbose output')
        PARSER.add_option ('-l', '--list', action='store_true', default=False, help='verbose output')
        PARSER.add_option ('-a', '--append', action='store', type='string', help='execute command after default commands')
        PARSER.add_option ('-c', '--commands', action='store', type='string', help='execute the command instead of default commands')
        PARSER.add_option ('-d', '--delimiter', action='store', type='string', default=',', help='delimiter between commands, default is ","')

        (OPTIONS, CLI_ARGS) = PARSER.parse_args()

        XTOOLS_HOSTS = parse_config_file()

        if OPTIONS.list:
            show_available_hosts(XTOOLS_HOSTS)
            sys.exit(0)

        if len(CLI_ARGS) < 1:
            PARSER.error ('missing argument')
        else:
            HOST_PATTERN = CLI_ARGS[0]

        main()

        sys.exit(0)
    except KeyboardInterrupt, e: # Ctrl-C
        sys.exit(1)
    except SystemExit, e: # sys.exit()
        raise e
    except Exception, e:
        print 'ERROR, UNEXPECTED EXCEPTION'
        print str(e)
        traceback.print_exc()
        os._exit(1)

# vi:ts=4:sw=4:expandtab:ft=python:
